#ifndef _AVLTREE_H_
#define _AVLTREE_H_

/*****************************************************************************/
/**********					AVLTree class declaration				**********/
/*****************************************************************************/

#include "Common.h"
#include "ICompare.h"
#include "IAVLTreeOwner.h"
#include "TreeNode.h"
#include "Exceptions.h"
#include "Stack.h"
#include <math.h>
#include <cassert>
#include <algorithm>

namespace DataStructuresWet1
{

template <class DataT>
class AVLTree
{
	/*************************************************/
	/**********			Public members			******/
	/*************************************************/
public:
	//Ctors
	AVLTree(ICompare<DataT>& comparator, IAVLTreeOwner<DataT>* owner = NULL);
	AVLTree(ICompare<DataT>& comparator, int size, IAVLTreeOwner<DataT>* owner = NULL);
	//Copy ctor
	AVLTree(const AVLTree& other);
	//Dtor
	~AVLTree();
	//Dictionary interface
	void Insert(DataT* data);
	void Remove(DataT& data);
	DataT* Find(DataT& data) const;
	//Rank tree interface
	DataT* Select(int index) const;
	//Merges 2 trees and returns the new merged tree
	AVLTree<DataT>* Merge(const AVLTree<DataT>& other,
		typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr = NULL);
	//Returns the size of the root element
	int Size() const;
	//Fills a sorted array from the tree data
	void ToSortedArray(DataT* arr[], typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr = NULL) const;
	//Fills an empty tree from a sorted array of DataT*
	void FromSortedArray(DataT* arr[], typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr = NULL);
	//Returns the minimal element in the tree (left most child of root)
	DataT* MinElement() const;
	//Returns the maximal element in the tree (right most child of root)
	DataT* MaxElement() const;
	//Returns the median element in the tree
	DataT* MedianElement() const;
	//Setters and getters for the owner of the tree (where the current instance resides)
	IAVLTreeOwner<DataT>* GetOwner() const;
	void SetOwner(IAVLTreeOwner<DataT>* value);
	//Empties the tree of all elements (by default the DataT* is not deleted)
	void Empty(bool deleteData = false);
	//Operator overloading
	AVLTree<DataT>& operator=(AVLTree<DataT>& other);
	/*************************************************/
	/**********			Private members			******/
	/*************************************************/
private:
	//internal typedefs
	typedef _TreeNode<DataT> TreeNode;
	//Finds the TreeNode conatining data
	TreeNode* FindNode(DataT& data) const;
	//Returns a new TreeNode* with a copy of the sub-tree with src as it's root
	TreeNode* CopySubTree(TreeNode* src);
	//Internal empty (recursive)
	void DoEmpty(TreeNode* node, bool deleteData);
	//Merges 2 sorted arrays, needed for the merge algorithms
	void MergeArrays(DataT* arr1[], int size1, DataT* arr2[], int size2, DataT* outArr[]);
	//ReBalances node if it is out of balance
	void ReBalance(TreeNode* node);
	//Internal select (recursive)
	DataT* DoSelect(TreeNode* node, int index) const;
	//Builds a full sub-tree under leaf root, of height height (recursive)
	void BuildFullSubTree(TreeNode* root, int height);
	//Removes numOfLeaves from a Full-Binary tree, starting from the right most leaf and continuing from there
	int RemoveLeavesFromFullTree(TreeNode* pos, int numOfLeaves);
	//Internal ToSortedArray (recursive)
	DataT** DoToSortedArray(DataT* arr[], TreeNode* node, typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr) const;
	//Internal FromSortedArray (recursive)
	DataT** DoFromSortedArray(DataT* arr[], TreeNode* node, typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr);
	//Data members
	ICompare<DataT> &Comparator;
	TreeNode* Root;
	IAVLTreeOwner<DataT>* Owner;
};

/*************************************************/
/**********		Class Implementation		******/
/*************************************************/
template <class DataT>
AVLTree<DataT>::AVLTree(ICompare<DataT> &comparator, IAVLTreeOwner<DataT>* owner) : Comparator(comparator), Root(NULL), Owner(owner) {}

template <class DataT>
AVLTree<DataT>::AVLTree(const AVLTree<DataT> &other) : Comparator(other.Comparator), Root(CopySubTree(other.Root)), Owner(other.Owner) {}

template <class DataT>
AVLTree<DataT>::AVLTree(ICompare<DataT> &comparator, int size, IAVLTreeOwner<DataT>* owner) : Comparator(comparator), Root(NULL), Owner(owner)
{
	if (size ==0)
	{
		return;
	}
	//Calculate size for full-tree
	int height = (int) ceil(log2((float) size+1));
	int diff = (int)(pow((float) 2, height)) - size - 1;
	Root = new TreeNode(NULL);
	//First build full tree of the minimal size bigger than the size we need
	BuildFullSubTree(Root, height - 1);
	//Remove leafs as needed
	RemoveLeavesFromFullTree(Root, diff);
	Root->UpdateHeight();
	Root->UpdateSize();
	assert(this->Size() == size);
}

template <class DataT>
AVLTree<DataT>::~AVLTree() 
{
	Empty();
}

template <class DataT>
void AVLTree<DataT>::Empty(bool deleteData)
{
	DoEmpty(this->Root, deleteData);
	this->Root = NULL;
}

template <class DataT>
void AVLTree<DataT>::DoEmpty(typename AVLTree<DataT>::TreeNode* node, bool deleteData)
{
	//Empty the tree by traversing postorder on the tree.
	//First delete left child and right child and only then delete the current node
	if (node == NULL)
	{
		return;
	}

	DoEmpty(node->GetLeftChild(), deleteData);
	DoEmpty(node->GetRightChild(), deleteData);

	node->MakeOrphan();
	if (deleteData)
	{
		delete node->GetDataPtr();
	}
	delete node;
}


template <class DataT>
typename AVLTree<DataT>::TreeNode* AVLTree<DataT>::CopySubTree(typename AVLTree<DataT>::TreeNode* src)
{
	//Copy the tree by traversing in preorder
	//First copy the current node, then left sub-tree and finally right sub-tree
	if (src == NULL)
	{
		return NULL;
	}
	TreeNode* node = new TreeNode(*src);
	node->SetLeftChild(CopySubTree(src->GetLeftChild()));
	node->SetRightChild(CopySubTree(src->GetRightChild()));
	return node;
}


template <class DataT>
AVLTree<DataT>& AVLTree<DataT>::operator=(AVLTree<DataT>& other)
{
	if (this != &other)
	{
		this->AVLTree<DataT>::~AVLTree(); //explicit non-virtual destructor
		new (this) AVLTree(other);
	}
	return *this;
}


template <class DataT>
DataT* AVLTree<DataT>::MedianElement() const
{
	//Simply select (ceil(size/2)-1)
	int medianIndex = (int) ceil((float) (this->Size()) / 2) - 1;
	return this->Select(medianIndex);
}

template <class DataT>
DataT* AVLTree<DataT>::MinElement() const
{
	//Left most child of root
	if (Root == NULL)
	{
		return NULL;
	}
	return Root->GetLeftMostChild()->GetDataPtr();
}

template <class DataT>
DataT* AVLTree<DataT>::MaxElement() const
{
	//Right most child of root
	if (Root == NULL)
	{
		return NULL;
	}
	return Root->GetRightMostChild()->GetDataPtr();
}

template <class DataT>
int AVLTree<DataT>::Size() const
{
	if (Root == NULL)
	{
		return 0;
	}
	return Root->GetSize();
}

template <class DataT>
void AVLTree<DataT>::BuildFullSubTree(typename AVLTree<DataT>::TreeNode* root, int height)
{
	if (height == 0)
	{
		return;
	}
	//First create current node
	TreeNode* leftChild = new TreeNode(NULL);
	//Build left sub-tree of height -1
	BuildFullSubTree(leftChild, height - 1);
	root->SetLeftChild(leftChild);
	//Build right sub-tree of height -1
	TreeNode* rightChild = new TreeNode(NULL);
	BuildFullSubTree(rightChild, height - 1);
	root->SetRightChild(rightChild);
}

template <class DataT>
int AVLTree<DataT>::RemoveLeavesFromFullTree(typename AVLTree<DataT>::TreeNode* pos, int numOfLeaves)
{
	//Removes numOfLeaves from the tree by traversing in reverse inorder,
	//and remove the current node only if it is a leaf
	if (pos == NULL)
	{
		return numOfLeaves;
	}
	numOfLeaves = RemoveLeavesFromFullTree(pos->GetRightChild(), numOfLeaves);
	pos->UpdateSize();
	pos->UpdateHeight();
	if (numOfLeaves == 0)
	{
		return numOfLeaves;
	}
	if (pos->NodeChildrenType() == _TreeNode<DataT>::LEAF)
	{
		pos->MakeOrphan();
		delete pos;
		numOfLeaves--;
		return numOfLeaves;
	}
	numOfLeaves = RemoveLeavesFromFullTree(pos->GetLeftChild(), numOfLeaves);
	pos->UpdateSize();
	pos->UpdateHeight();
	return numOfLeaves;
}


template <class DataT>
void AVLTree<DataT>::ToSortedArray(DataT* arr[], 
	typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr) const
{
	DoToSortedArray(arr, this->Root, funcPtr);
}

template <class DataT>
DataT** AVLTree<DataT>::DoToSortedArray(DataT* arr[], 
	typename AVLTree<DataT>::TreeNode* node, 
	typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr) const
{
	//Traverse the tree in inorder and fill arr.
	//Also calls the owner funcPtr if the owner needs to change the data as traversing
	if (node != NULL)
	{
		arr = DoToSortedArray(arr, node->GetLeftChild(), funcPtr);
		*arr = node->GetDataPtr();
		if (funcPtr != NULL)
		{
			(Owner->*funcPtr)(**arr);
		}
		arr = DoToSortedArray(++arr, node->GetRightChild(), funcPtr);
	}
	return arr;
}

template <class DataT>
DataT** AVLTree<DataT>::DoFromSortedArray(DataT* arr[], 
	typename AVLTree<DataT>::TreeNode* node, 
	typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr)
{
	//Traverse the tree in inorder and fill the empty tree from arr
	//Also calls the owner funcPtr if the owner needs to change the data as traversing
	if (node != NULL)
	{
		arr = DoFromSortedArray(arr, node->GetLeftChild(), funcPtr);
		if (funcPtr != NULL)
		{
			(Owner->*funcPtr)(**arr);
		}
		node->SetDataPtr(*arr);
		arr = DoFromSortedArray(++arr, node->GetRightChild(), funcPtr);
	}
	return arr;
}


template <class DataT>
void AVLTree<DataT>::FromSortedArray(DataT* arr[], 
	typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr)
{
	DoFromSortedArray(arr, this->Root, funcPtr);
}

template <class DataT>
DataT* AVLTree<DataT>::Find(DataT& data) const
{
	//Simply calls FindNode
	TreeNode *node = FindNode(data);
	if (node == NULL)
	{
		return NULL;
	}
	return node->GetDataPtr();
}

template <class DataT>
typename  AVLTree<DataT>::TreeNode* AVLTree<DataT>::FindNode(DataT& data) const
{
	TreeNode *node = this->Root;
	while(node != NULL)
	{
		//If the data is in node, then we are done
		if (*(node->GetDataPtr()) == data)
		{
			break;
		}
		//If it is less, visit left sub-tree
		if (Comparator.IsLess(data, *(node->GetDataPtr())))
		{
			node = node->GetLeftChild();
		}
		//If it is more, visit right sub-tree
		else
		{

			node = node->GetRightChild();
		}
		
	}
	//node not found
	if (node == NULL)
	{
		return NULL;
	}
	return node;
}

template <class DataT>
void AVLTree<DataT>::Insert(DataT* data)
{
	TreeNode* node = this->Root;
	if (node == NULL)
	{
		this->Root = new TreeNode(data);
		return;
	}
	//Find  insert position
	while(node != NULL)
	{
		//if data is already in tree, throw an exception
		if (*node->GetDataPtr() == *data)
		{
			throw FailureException();
		}
		if (Comparator.IsLess(*data, *(node->GetDataPtr())))
		{
			if (node->GetLeftChild() == NULL)
			{
				node->SetLeftChild(new TreeNode(data));
				break;
			}
			else
			{
				node = node->GetLeftChild();
			}
		}
		else
		{
			if (node->GetRightChild() == NULL)
			{
				node->SetRightChild(new TreeNode(data));
				break;
			}
			else
			{
				node = node->GetRightChild();
			}
		}
	}
	assert(node != NULL);
	while(node != NULL)
	{
		ReBalance(node);
		node = node->GetParent();
	}
}


template <class DataT>
void AVLTree<DataT>::Remove(DataT& data)
{
	//If node is not in tree, throw an exception
	TreeNode* node = this->FindNode(data);
	if (node == NULL)
	{
		throw FailureException();
	}
	TreeNode* ReBalanceFrom;
	//Replace node with a proper substitute to make it a leaf (viable for removal)
	switch(node->NodeChildrenType())
	{
	case _TreeNode<DataT>::LEAF:
		ReBalanceFrom = node->GetParent();
		break;
	case _TreeNode<DataT>::HAS_LEFT_CHILD:
		ReBalanceFrom = node->GetParent();
		if (node == Root)
		{
			Root = node->GetLeftChild();
		}
		node->GetLeftChild()->ReplaceParent();
		break;
	case _TreeNode<DataT>::HAS_RIGHT_CHILD:
		ReBalanceFrom = node->GetParent();
		if (node == Root)
		{
			Root = node->GetRightChild();
		}
		node->GetRightChild()->ReplaceParent();
		break;
	case _TreeNode<DataT>::HAS_TWO_CHILDREN:
		node->SwitchPosition(*(node->GetRightChild()->GetLeftMostChild()));
		node = node->GetRightChild()->GetLeftMostChild();
		ReBalanceFrom = node->GetParent();
		if (node->NodeChildrenType() == _TreeNode<DataT>::HAS_RIGHT_CHILD)
		{
			if (node == Root)
			{
				Root = node->GetRightChild();
			}
			node->GetRightChild()->ReplaceParent();
		}
		break;
	}
	//Now that node is leaf, delete it
	assert(node->NodeChildrenType() == _TreeNode<DataT>::LEAF);
	node->MakeOrphan();
	if (node == Root)
	{
		Root = NULL;
	}
	delete node;
	node = ReBalanceFrom;
	//Re-balance tree on search path
	while(node != NULL)
	{
		ReBalance(node);
		node = node->GetParent();
	}
}

template <class DataT>
AVLTree<DataT>* AVLTree<DataT>::Merge(const AVLTree<DataT>& other,
		typename IAVLTreeOwner<DataT>::DoActionOnDataFunctionPtr funcPtr)
{
	int size1 = this->Size();
	int size2 = other.Size();
	DataT** arr1;
	DataT** arr2;
	DataT** newArr;
	if (size1 == 0)
	{
		arr1 = NULL;
	}
	else
	{
		arr1 = new DataT*[size1];
	}
	if (size2 == 0)
	{
		arr2 = NULL;
	}
	else
	{
		try
		{
			arr2 = new DataT*[size2];
		}
		catch (std::bad_alloc ex)
		{
			delete [] arr1;
			throw ex;
		}
	}
	//Send this tree to sorted array
	this->ToSortedArray(arr1);
	//Send other tree to sorted array
	other.ToSortedArray(arr2, funcPtr);
	int newSize = size1 + size2;
	if (newSize == 0)
	{
		newArr = NULL;
	}
	else
	{
		try
		{
			newArr = new DataT*[newSize];
		}
		catch (std::bad_alloc ex)
		{
			delete [] arr1;
			delete [] arr2;
			throw ex;
		}
	}
	//Merge the 2 arrays
	MergeArrays(arr1, size1, arr2, size2, newArr);
	AVLTree<DataT> *newTree;
	//Create new empty tree of size newSize
	try
	{
		newTree = new AVLTree<DataT>(this->Comparator, newSize, this->Owner);
	}
	catch (std::bad_alloc ex)
	{
		delete [] arr1;
		delete [] arr2;
		delete [] newArr;
		throw ex;
	}
	//Fill the empty tree from the merged sorted array
	newTree->FromSortedArray(newArr);
	//Delete the temporary arrays
	delete [] arr1;
	delete [] arr2;
	delete [] newArr;
	return newTree;
}

template <class DataT>
void AVLTree<DataT>::ReBalance(TreeNode* node)
{
	node->UpdateSize();
	node->UpdateHeight();
	//node->UpdateWeight();
	int nodeBF = node->GetBalanceFactor();
	int childBF;
	assert(abs(nodeBF) <= 2);
	TreeNode* newRoot = Root;
	if (nodeBF == 2)
	{
		childBF = node->GetLeftChild()->GetBalanceFactor();
		assert(abs(childBF) <= 2);
		//rotate LL
		if (childBF >= 0)
		{
			if (node == Root)
			{
				newRoot = Root->GetLeftChild();
			}
			node->RotateLL();
		}
		//rotate LR
		else //childBF == -1
		{
			if (node == Root)
			{
				newRoot = Root->GetLeftChild()->GetRightChild();
			}
			node->RotateLR();
		}
	}
	if (nodeBF == -2)
	{
		childBF = node->GetRightChild()->GetBalanceFactor();
		assert(abs(childBF) <= 2);
		//rotate RR
		if (childBF <= 0)
		{
			if (node == Root)
			{
				newRoot = Root->GetRightChild();
			}
			node->RotateRR();
		}
		//rotate RL
		else //childBF == 1
		{
			if (node == Root)
			{
				newRoot = Root->GetRightChild()->GetLeftChild();
			}
			node->RotateRL();
		}
	}
	Root = newRoot;
}

template <class DataT>
DataT* AVLTree<DataT>::Select(int index) const
{
	TreeNode* root = this->Root;
	return DoSelect(root, index);
}

template <class DataT>
DataT* AVLTree<DataT>::DoSelect(typename AVLTree<DataT>::TreeNode* node, int index) const
{
	//index not found:
	if (node == NULL)	
	{
		return NULL;
	}
	//index found, return data
	if (node->GetWeight() == index)
	{
		return node->GetDataPtr();
	}
	//If index is less, search left sub-tree
	if (node->GetWeight() > index)
	{
		return DoSelect(node->GetLeftChild(), index);
	}
	//If index is more, search right sub-tree for index-weight of current tree - 1
	else //node->GetWeight() < index
	{
		return DoSelect(node->GetRightChild(), index - node->GetWeight() - 1);
	}
}

template <class DataT>
void AVLTree<DataT>::MergeArrays(DataT* arr1[], int size1, DataT* arr2[], int size2, DataT* outArr[])
{
      int i = 0;
      int j = 0;
      int k = 0;
      while (i < size1 && j < size2) 
	  {
            if (Comparator.IsLeq(*(arr1[i]), *(arr2[j])))	//arr1[i] <= arr2[j]
			{
                  outArr[k] = arr1[i];
                  i++;
            } else											//arr1[i] > arr2[j]
			{
                  outArr[k] = arr2[j];
                  j++;
            }
            k++;
      }
      if (i < size1) 
	  {
            for (int p = i; p < size1; p++) 
			{
                  outArr[k] = arr1[p];
                  k++;
            }
      } else 
	  {
            for (int p = j; p < size2; p++) 
			{
                  outArr[k] = arr2[p];
                  k++;
            }
      }
}

template <class DataT>
IAVLTreeOwner<DataT>* AVLTree<DataT>::GetOwner() const
{
	return this->Owner;
}
template <class DataT>

void AVLTree<DataT>::SetOwner(IAVLTreeOwner<DataT>* value)
{
	this->Owner = value;
}

} //DataStructuresWet1


#endif //_AVLTREE_H_
